package io.github.javpower.knowledgerag.service;

import com.alibaba.fastjson2.JSON;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import io.github.javpower.knowledgerag.config.RagConfig;
import io.github.javpower.knowledgerag.core.ContextBuilder;
import io.github.javpower.knowledgerag.core.DocumentReader;
import io.github.javpower.knowledgerag.core.Vectorizer;
import io.github.javpower.knowledgerag.req.QuestionReq;
import io.github.javpower.vectorexclient.VectorRexClient;
import io.github.javpower.vectorexclient.builder.QueryBuilder;
import io.github.javpower.vectorexclient.entity.MetricType;
import io.github.javpower.vectorexclient.entity.ScalarField;
import io.github.javpower.vectorexclient.entity.VectorFiled;
import io.github.javpower.vectorexclient.req.CollectionDataAddReq;
import io.github.javpower.vectorexclient.req.VectoRexCollectionReq;
import io.github.javpower.vectorexclient.res.PageResult;
import io.github.javpower.vectorexclient.res.ServerResponse;
import io.github.javpower.vectorexclient.res.VectorSearchResult;
import io.github.javpower.vectorexclient.util.GsonUtil;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.util.*;
import java.util.stream.Collectors;

@Service
public class QnAService {

    private String collectionName = "";
    private VectorRexClient client = null;
    private String OLLAMA_API_URL = "";
    private String document = "";
    private String model = "";

    private static final Gson gson = new Gson();


    @Autowired
    public QnAService(RagConfig config) {
//        collectionName=config.getCollectionName();
//        OLLAMA_API_URL=config.getOllamaApiUrl();
//        document=config.getDocument();
//        model=config.getModel();
//        this.client =  new VectorRexClient(config.getVectorUrl(), config.getVectorUser(), config.getVectorPassword());
//        initialize();
        try {
            collectionName=config.getCollectionName();
            OLLAMA_API_URL=config.getOllamaApiUrl();
            document=config.getDocument();
            model=config.getModel();
            this.client =  new VectorRexClient(config.getVectorUrl(), config.getVectorUser(), config.getVectorPassword());
            initialize();
        } catch (Exception e) {
            e.printStackTrace();
        }


    }
    private void initialize() {
        boolean anyMatch = client.getCollections().getData().stream().anyMatch(collection -> collection.getCollectionName().equals(collectionName));
        if (anyMatch) {
            return;
        }
        List<String> chunks = new ArrayList<>();
        List<List<Float>> vectors = new ArrayList<>();

        // 读取文档并分块
        chunks.addAll(DocumentReader.readAndChunkDocument(document, 1024));

        // 向量化
        vectors.addAll(Vectorizer.vectorizeChunks(OLLAMA_API_URL,model,chunks));

        // 创建集合
        createCollection(collectionName, vectors.get(0).size());

        // 添加数据
        for (int i = 0; i < chunks.size(); i++) {
            addData(collectionName,i,chunks.get(i), vectors.get(i));
            String data = "collectionName: " + collectionName + ",id:  " + i + ",text: " + chunks.get(i) + ",vector: " + vectors.get(i);
            System.out.println(data);
            for (Float result : vectors.get(i)) {
                System.out.println("vectors【"+i+"】: " + result);
            }
        }
    }

    private void createCollection(String collectionName, int vectorDimensions) {

        List<ScalarField> scalarFields = new ArrayList<>();
        scalarFields.add(ScalarField.builder().name("id").isPrimaryKey(true).build());
        scalarFields.add(ScalarField.builder().name("text").isPrimaryKey(false).build());

        List<VectorFiled> vectorFields = new ArrayList<>();
        vectorFields.add(VectorFiled.builder().name("vector").metricType(MetricType.FLOAT_COSINE_DISTANCE).dimensions(vectorDimensions).build());

        VectoRexCollectionReq req = VectoRexCollectionReq.builder()
                .collectionName(collectionName)
                .scalarFields(scalarFields)
                .vectorFileds(vectorFields)
                .build();

        ServerResponse<Void> response = client.createCollection(req);
        if (response.isSuccess()) {
            System.out.println("集合创建成功");
        } else {
            System.out.println("集合创建失败：" + response.getMsg());
        }
    }

    private void addData(String collectionName, Integer id, String text, List<Float> vector) {
        Map<String, Object> metadata = new HashMap<>();
        metadata.put("id", id);
        metadata.put("text", text);
        metadata.put("vector", vector);

        CollectionDataAddReq req = CollectionDataAddReq.builder()
                .collectionName(collectionName)
                .metadata(metadata)
                .build();

        ServerResponse<Void> response = client.addCollectionData(req);
        if (response.isSuccess()) {
            System.out.println("数据添加成功");
        } else {
            System.out.println("数据添加失败：" + response.getMsg());
        }
    }

    private List<VectorSearchResult> queryVectorDatabase(String collectionName, List<Float> queryVector, int topK) {
        QueryBuilder builder = QueryBuilder.lambda(collectionName);
        builder.vector("vector", queryVector).topK(topK);
        ServerResponse<List<VectorSearchResult>> response = client.queryCollectionData(builder);
        if (response.isSuccess()) {
            return response.getData();
        } else {
            System.out.println("数据查询失败：" + response.getMsg());
            return Collections.emptyList();
        }
    }

    public String generateAnswer(QuestionReq req) {
        try {
            // 将问题转换为向量
            List<Float> queryVector = Vectorizer.getVector(OLLAMA_API_URL,model,req.getQuestion());

            // 查询向量数据库
            List<VectorSearchResult> results = queryVectorDatabase(collectionName, queryVector, 2);

            // 构建上下文
            String context = ContextBuilder.buildContext(results);

            // 使用模型生成答案
            String answer = executeQwenModel(model,req.getQuestion(),context);

            return answer;
        } catch (Exception e) {
            e.printStackTrace();
            return "Error generating answer";
        }
    }

    private String executeQwenModel(String model,String question, String context) {
        OkHttpClient client = new OkHttpClient();
        MediaType mediaType = MediaType.parse("application/json");
        String json = "{\"model\": \""+model+"\", \"prompt\": \"" + question + " " + context + "\"}";
        RequestBody body = RequestBody.create(json, mediaType);
        Request request = new Request.Builder()
                .url(OLLAMA_API_URL+"/generate")
                .post(body)
                .addHeader("Content-Type", "application/json")
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("Unexpected code " + response);
            }
            String res = response.body().string();
            String[] split = res.split("\n");
            String answerRes = Arrays.stream(split).map(v-> {
                Map s = GsonUtil.fromJson(v, Map.class);
                return s.get("response").toString();
            }).collect(Collectors.joining(""));
            return answerRes;
        } catch (IOException e) {
            e.printStackTrace();
            return "Error generating answer";
        }
    }

    public String vectorQuery(String queryString) {
        QueryBuilder builder = QueryBuilder.lambda(collectionName);
        builder.like("text", "vector")
                .vector("vector", Arrays.asList(0.1f, 0.2f, 0.3f)) // 向量查询条件
                .topK(10); // 返回最相似的 10 条结果
        //ServerResponse<List<VectorSearchResult>> response = client.queryCollectionData(builder);
        builder.page(1, 10); // 查询第 1 页，每页 10 条数据
        ServerResponse<PageResult<VectorSearchResult>> queryResponse = client.pageCollectionData(builder);
        if (queryResponse.isSuccess()) {
            List<VectorSearchResult> results = (List<VectorSearchResult>) queryResponse.getData();
            for (VectorSearchResult result : results) {
                System.out.println("查询结果：" + result.getId() + ", 分数：" + result.getScore());
            }
            return GsonUtil.toJson(queryResponse.getData());
        } else {
            System.out.println("数据查询失败：" + queryResponse.getMsg());
        }
        return null;
    }







    public String functionCall(String queryString) {
        String userInput = "我想知道北京的天气";
        String prompt = constructPrompt(userInput);
        String response = executeModel(model, prompt);
        JsonObject functionCall = parseFunctionCall(response);
        String result = "";
        if (functionCall!= null) {
            result = callFunction(functionCall);
            System.out.println(result);
        } else {
            System.out.println(response);
        }
        return result;
    }


    // 定义函数列表
    private final Map<String, Map<String, Object>> functions = new HashMap<>() {{
        put("get_weather", Map.of(
                "description", "获取指定城市的天气",
                "parameters", Map.of(
                        "city", Map.of(
                                "type", "string",
                                "description", "城市名称"
                        )
                )
        ));
    }};


    private String constructPrompt(String userInput) {
        return "你现在需要根据用户需求调用函数获取信息。用户需求：" + userInput +
                "。可调用函数：" + JSON.toJSONString(functions);
    }



    private String executeModel(String model,String prompt) {
        OkHttpClient client = new OkHttpClient();
        MediaType mediaType = MediaType.parse("application/json");
        String json = "{\"model\": \""+model+"\", \"prompt\": \"" + prompt + "\"}";
        RequestBody body = RequestBody.create(json, mediaType);
        Request request = new Request.Builder()
                .url(OLLAMA_API_URL+"/generate")
                .post(body)
                .addHeader("Content-Type", "application/json")
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("Unexpected code " + response);
            }
            String res = response.body().string();
            String[] split = res.split("\n");
            String answerRes = Arrays.stream(split).map(v-> {
                Map s = GsonUtil.fromJson(v, Map.class);
                return s.get("response").toString();
            }).collect(Collectors.joining(""));
            return answerRes;
        } catch (IOException e) {
            e.printStackTrace();
            return "Error generating answer";
        }
    }

    private String callOllama(String prompt) {
        try {
            URL url = new URL(OLLAMA_API_URL + "/generate");
            HttpURLConnection connection = (HttpURLConnection) url.openConnection();
            connection.setRequestMethod("POST");
            connection.setRequestProperty("Content-Type", "application/json");
            connection.setDoOutput(true);

            JsonObject requestBody = new JsonObject();
            requestBody.addProperty("model", model);
            requestBody.addProperty("prompt", prompt);

            try (OutputStream os = connection.getOutputStream()) {
                byte[] input = requestBody.toString().getBytes("utf-8");
                os.write(input, 0, input.length);
            }

            try (BufferedReader br = new BufferedReader(
                    new InputStreamReader(connection.getInputStream(), "utf-8"))) {
                StringBuilder response = new StringBuilder();
                String responseLine;
                while ((responseLine = br.readLine())!= null) {
                    response.append(responseLine.trim());
                }

                String res = response.toString();
                String[] split = res.split("\n");
                String answerRes = Arrays.stream(split).map(v-> {
                    Map s = GsonUtil.fromJson(v, Map.class);
                    return s.get("response").toString();
                }).collect(Collectors.joining(""));

                //JsonObject jsonResponse = JsonParser.parseString(response.toString()).getAsJsonObject();
                return answerRes;
            }
        } catch (IOException e) {
            e.printStackTrace();
            return "";
        }
    }

    private static JsonObject parseFunctionCall(String response) {
        try {
            return JsonParser.parseString(response).getAsJsonObject();
        } catch (Exception e) {
            return null;
        }
    }

    private static String callFunction(JsonObject functionCall) {
        String functionName = functionCall.get("function_name").getAsString();
        JsonObject parameters = functionCall.get("parameters").getAsJsonObject();

        // 这里是模拟真实函数调用，实际应用中需要实现真实的函数逻辑
        if ("get_weather".equals(functionName)) {
            String city = parameters.get("city").getAsString();
            return city + "的天气是晴天";
        }
        return "未找到对应的函数实现";
    }



    private static String getWeather(String city) {
        // 这里假设使用OpenWeatherMap API来获取天气，实际需要替换为真实的API密钥
//        String apiKey = "YOUR_API_KEY";
//        String weatherApiUrl = "http://api.openweathermap.org/data/2.5/weather?q=" + city + "&appid=" + apiKey + "&units=metric";
//        try {
//            URL url = new URL(weatherApiUrl);
//            HttpURLConnection connection = (HttpURLConnection) url.openConnection();
//            connection.setRequestMethod("GET");
//
//            try (BufferedReader br = new BufferedReader(
//                    new InputStreamReader(connection.getInputStream(), "utf-8"))) {
//                StringBuilder response = new StringBuilder();
//                String responseLine;
//                while ((responseLine = br.readLine())!= null) {
//                    response.append(responseLine.trim());
//                }
//                JsonObject jsonResponse = JsonParser.parseString(response.toString()).getAsJsonObject();
//                // 假设天气信息在main -> temp字段中，实际根据API文档调整
//                double temperature = jsonResponse.getAsJsonObject("main").get("temp").getAsDouble();
//                return city + "的当前温度是 " + temperature + "°C";
//            }
//        } catch (IOException e) {
//            e.printStackTrace();
//            return "获取天气信息失败";
//        }
        return city + "的当前温度是 " + 666 + "°C";
    }
}